from mdp.nodes import GrowingNeuralGasNode
from itertools import izip
import numpy,random
from scipy import interpolate
import sys
sys.setrecursionlimit(100000)

class Predictor:
    def __init__(self, planner, actionChoices, numSensors,drawer = None):
        self.numSensors = numSensors
        self.actionChoices = actionChoices
        self.planner = planner
        self.history = []
        self.info = {action: list() for action in self.actionChoices}
        self.gngs = dict((action,GrowingNeuralGasNode(#eps_b = .5, lambda_=20,max_age=15, eps_n= .2
                                                      )) for action in self.actionChoices)
        self.interpolators = {}
        self.possActions = {}
        self.actionComps = {action:self.gngs[action].graph.connected_components() for action in self.actionChoices}
        print "ready"
        self.previousMove = None
        self.drawer = drawer
        self.turn = 0
        
        
    def save(self):
        import cPickle
        f = open('ai.pickle','w')
        cPickle.dump(self.gngs,f)
        f.close()
        print "Saved"
        return 1
    
    def load(self):
        import cPickle
        f = open('ai.pickle','r')
        self.gngs = cPickle.load(f)
        print "Loaded"
        return 1
            
        
    def tick(self, action, sensors):
        # prevstate ----[prev]---> sensors ----[action]---->
        self.history.append(sensors)
        if sensors == (0,0,0,0) and len(self.history) > 1:
            print "WTF"
            raw_input()
        if len(self.history) > 1:
            self.train(self.previousMove,self.history[-2],self.history[-1])
            self.info[self.previousMove].append(self.history[-2]+self.history[-1])
        
        self.previousMove = action
        self.turn += 1
        return 1
    
    
    def testDrive(self, *sensors):
        return self.drive(sensors)
        
    def smartDrive(self, *sensors):
        return self.drive(sensors,True)
    
    def drive(self, sensors, smart=False):
        #=0  - move left
        #=1  - move right
        
        if self.drawer is not None:
            self.drawer.draw(sensors,self.turn)
                
        if smart:
            nextMove = self.planner.planner(sensors, self.actionChoices, self.possActionResults,self.turn)
        else:
            nextMove = None
        
        if nextMove is None:
            nextMove = random.choice(self.actionChoices)
        
        self.tick(nextMove,sensors)
        return nextMove
    

    def translate_state_via_comp(self, state, comp):
        if tuple(comp) in self.interpolators:
            interps = self.interpolators[tuple(comp)]
        else:
            points = []
            values = []    
            for node in comp:
                points.append(node.data.pos[:self.numSensors])
                values.append(numpy.array(node.data.pos[self.numSensors:]))
            
            
            interps = [interpolate.Rbf(*(map(lambda x:numpy.array(x),zip(*points))+[numpy.array([v[s] for v in values])]), smooth=.0005) for s in range(self.numSensors)] 
            self.interpolators[tuple(comp)] = interps
         
        newState = tuple([float(interps[s](*state)) for s in range(self.numSensors)])
        
        if newState == (0,0,0,0):
            print "WAT"
            raw_input()
        
        return newState 
        
    def state_near_comp_beginning(self, state, comp):   
        if comp not in self.compBeginningArrays:
            self.compBeginningArrays[comp] = [numpy.array(z) for z in izip(*[list(node.data.pos[:self.numSensors]) for node in comp])]
            
        nodeArrays = self.compBeginningArrays[comp]

        #Try each possible start node
        x = all(
                map(
                    lambda (nodeD, stateD): 
                        abs((stateD - numpy.mean(nodeD))/numpy.std(nodeD)) 
                            <= 5,
                    izip(nodeArrays,state)))
        

        return x

    def possActionResults(self, state, action):
        if (state,action) not in self.possActions:
            actionComps = self.actionComps[action]
            self.possActions[(state,action)] = []
            for comp in actionComps:
                if len(comp) < 6:
                    continue
                comp = tuple(comp)
                if self.state_near_comp_beginning(state,comp):
                    poss =self.translate_state_via_comp(state, comp)
                    self.possActions[(state,action)].append(poss) 
                   
        return tuple(self.possActions[(state,action)])    

    
    def clearCache(self):
        self.possActions = {}
        self.compBeginningArrays = {}
        self.interpolators = {}
        

    def train(self,action,stateFrom,stateTo):
        pred = tuple(self.possActionResults(stateFrom, action))
        if pred == () or numpy.linalg.norm(numpy.array(pred) - numpy.array(stateTo)) > 1e-1:
            self.clearCache()
            self.gngs[action].train(numpy.array([stateFrom+stateTo]))
            self.actionComps = {action:self.gngs[action].graph.connected_components() for action in self.actionChoices}
        
    def reset(self):
        print  "last position, guess/actual outcome:"
        print self.previousMove
        if len(self.info[self.previousMove]) == 0:
            # Failed on first use of previousmove, couldn't collect last position..
            lastMove = (self.history[-1])+("UNKNOWN",)*self.numSensors
        else:
            lastMove = self.info[self.previousMove][-1]
        pred = self.possActionResults(lastMove[:self.numSensors],self.previousMove)
        print tuple(lastMove[:self.numSensors]), "\n", tuple(pred), "\n", tuple(lastMove[self.numSensors:])
        
        #=======================================================================
        # l = 0
        # for action in self.actionChoices:
        #    if len(self.info[action]) > 1:
        #        for s in self.info[action]:
        #            pred = tuple(self.possActionResults(s[:self.numSensors], action))
        #            if pred == () or numpy.linalg.norm(numpy.array(pred) - numpy.array(s[self.numSensors:])) > 1e-1:
        #                self.clearCache()
        #                self.gngs[action].train(numpy.array([s]))
        #                self.actionComps = {action:self.gngs[action].graph.connected_components() for action in self.actionChoices}
        #                l+=1
        #print "learned %s facts" % (l,)
        #=======================================================================         
                
        self.info = {k: list() for k in self.actionChoices}
        self.history = []
        self.planner.reset()
        
        self.previousMove = None
        self.turn = 0
        return 1